import json
import os.path


def get_data(data_type, callback=None):
    assert data_type in ['train', 'dev', 'test']
    dir_name = os.path.dirname(__file__)
    input_file = os.path.join(dir_name, 'data', data_type + '.json')
    features = []
    with open(input_file, 'r', encoding="utf-8") as f:
        for line in f:
            line = json.loads(line.strip())
            text = line['text']
            label_entities = line.get('label', None)
            line = {"text": text, "labels": label_entities}
            if callback is not None:
                line = callback(line=line)
            features.append(line)
    return features


def get_file_name(mode):
    assert mode in ['train', 'dev', 'test']
    dir_name = os.path.dirname(__file__)
    input_file = os.path.join(dir_name, 'data', mode + '.json')
    return input_file
